{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Train simple audio recognition model",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pO4-CY_TCZZS",
        "colab_type": "text"
      },
      "source": [
        "# Train a Simple Audio Recognition model for microcontroller use"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BaFfr7DHRmGF",
        "colab_type": "text"
      },
      "source": [
        "This notebook demonstrates how to train a 20kb [Simple Audio Recognition](https://www.tensorflow.org/tutorials/sequences/audio_recognition) model for [TensorFlow Lite for Microcontrollers](https://tensorflow.org/lite/microcontrollers/overview). It will produce the same model used in the [micro_speech](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/micro/examples/micro_speech) example application.\n",
        "\n",
        "The model is designed to be used with [Google Colaboratory](https://colab.research.google.com).\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/train_speech_model.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/train_speech_model.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XaVtYN4nlCft",
        "colab_type": "text"
      },
      "source": [
        "The notebook runs Python scripts to train and freeze the model, and uses the TensorFlow Lite converter to convert it for use with TensorFlow Lite for Microcontrollers.\n",
        "\n",
        "**Training is much faster using GPU acceleration.** Before you proceed, ensure you are using a GPU runtime by going to **Runtime -> Change runtime type** and selecting **GPU**. Training 18,000 iterations will take 1.5-2 hours on a GPU runtime.\n",
        "\n",
        "## Configure training\n",
        "\n",
        "The following `os.environ` lines can be customized to set the words that will be trained for, and the steps and learning rate of the training. The default values will result in the same model that is used in the micro_speech example. Run the cell to set the configuration:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ludfxbNIaegy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "\n",
        "# A comma-delimited list of the words you want to train for.\n",
        "# The options are: yes,no,up,down,left,right,on,off,stop,go\n",
        "# All other words will be used to train an \"unknown\" category.\n",
        "os.environ[\"WANTED_WORDS\"] = \"yes,no\"\n",
        "\n",
        "# The number of steps and learning rates can be specified as comma-separated\n",
        "# lists to define the rate at each stage. For example,\n",
        "# TRAINING_STEPS=15000,3000 and LEARNING_RATE=0.001,0.0001\n",
        "# will run 18,000 training loops in total, with a rate of 0.001 for the first\n",
        "# 15,000, and 0.0001 for the final 3,000.\n",
        "os.environ[\"TRAINING_STEPS\"]=\"15000,3000\"\n",
        "os.environ[\"LEARNING_RATE\"]=\"0.001,0.0001\"\n",
        "\n",
        "# Calculate the total number of steps, which is used to identify the checkpoint\n",
        "# file name.\n",
        "total_steps = sum(map(lambda string: int(string),\n",
        "                  os.environ[\"TRAINING_STEPS\"].split(\",\")))\n",
        "os.environ[\"TOTAL_STEPS\"] = str(total_steps)\n",
        "\n",
        "# Print the configuration to confirm it\n",
        "!echo \"Training these words: ${WANTED_WORDS}\"\n",
        "!echo \"Training steps in each stage: ${TRAINING_STEPS}\"\n",
        "!echo \"Learning rate in each stage: ${LEARNING_RATE}\"\n",
        "!echo \"Total number of training steps: ${TOTAL_STEPS}\"\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCgeOpvY9pAi",
        "colab_type": "text"
      },
      "source": [
        "## Install dependencies\n",
        "\n",
        "Next, we'll install a GPU build of TensorFlow, so we can use GPU acceleration for training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nd1iM1o2ymvA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Replace Colab's default TensorFlow install with a more recent\n",
        "# build that contains the operations that are needed for training\n",
        "!pip uninstall -y tensorflow tensorflow_estimator tensorboard\n",
        "!pip install -q tf-estimator-nightly==1.14.0.dev2019072901 tf-nightly-gpu==1.15.0.dev20190729"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T9Ty5mR58E4i",
        "colab_type": "text"
      },
      "source": [
        "We'll also clone the TensorFlow repository, which contains the scripts that train and freeze the model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "APGx0fEh7hFF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Clone the repository from GitHub\n",
        "!git clone -q https://github.com/tensorflow/tensorflow\n",
        "# Check out a commit that has been tested to work\n",
        "# with the build of TensorFlow we're using\n",
        "!git -c advice.detachedHead=false -C tensorflow checkout 17ce384df70"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aV_0qkYh98LD",
        "colab_type": "text"
      },
      "source": [
        "## Load TensorBoard\n",
        "\n",
        "Now, set up TensorBoard so that we can graph our accuracy and loss as training proceeds."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yZArmzT85SLq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Delete any old logs from previous runs\n",
        "!rm -rf /content/retrain_logs\n",
        "# Load TensorBoard\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir /content/retrain_logs"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x1J96Ron-O4R",
        "colab_type": "text"
      },
      "source": [
        "## Begin training\n",
        "\n",
        "Next, run the following script to begin training. The script will first download the training data:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VJsEZx6lynbY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python tensorflow/tensorflow/examples/speech_commands/train.py \\\n",
        "--model_architecture=tiny_conv --window_stride=20 --preprocess=micro \\\n",
        "--wanted_words=${WANTED_WORDS} --silence_percentage=25 --unknown_percentage=25 \\\n",
        "--quantize=1 --verbosity=WARN --how_many_training_steps=${TRAINING_STEPS} \\\n",
        "--learning_rate=${LEARNING_RATE} --summaries_dir=/content/retrain_logs \\\n",
        "--data_dir=/content/speech_dataset --train_dir=/content/speech_commands_train \\\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XQUJLrdS-ftl",
        "colab_type": "text"
      },
      "source": [
        "## Freeze the graph\n",
        "\n",
        "Once training is complete, run the following cell to freeze the graph."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xyc3_eLh9sAg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python tensorflow/tensorflow/examples/speech_commands/freeze.py \\\n",
        "--model_architecture=tiny_conv --window_stride=20 --preprocess=micro \\\n",
        "--wanted_words=${WANTED_WORDS} --quantize=1 --output_file=/content/tiny_conv.pb \\\n",
        "--start_checkpoint=/content/speech_commands_train/tiny_conv.ckpt-${TOTAL_STEPS}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_DBGDxVI-nKG",
        "colab_type": "text"
      },
      "source": [
        "## Convert the model\n",
        "\n",
        "Run this cell to use the TensorFlow Lite converter to convert the frozen graph into the TensorFlow Lite format, fully quantized for use with embedded devices."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lBj_AyCh1cC0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!toco \\\n",
        "--graph_def_file=/content/tiny_conv.pb --output_file=/content/tiny_conv.tflite \\\n",
        "--input_shapes=1,49,40,1 --input_arrays=Reshape_2 --output_arrays='labels_softmax' \\\n",
        "--inference_type=QUANTIZED_UINT8 --mean_values=0 --std_dev_values=9.8077"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dt6Zqbxu-wIi",
        "colab_type": "text"
      },
      "source": [
        "The following cell will print the model size, which will be under 20 kilobytes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XohZOTjR8ZyE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "model_size = os.path.getsize(\"/content/tiny_conv.tflite\")\n",
        "print(\"Model is %d bytes\" % model_size)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2pQnN0i_-0L2",
        "colab_type": "text"
      },
      "source": [
        "Finally, we use xxd to transform the model into a source file that can be included in a C++ project and loaded by TensorFlow Lite for Microcontrollers."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eoYyh0VU8pca",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Install xxd if it is not available\n",
        "!apt-get -qq install xxd\n",
        "# Save the file as a C source file\n",
        "!xxd -i /content/tiny_conv.tflite > /content/tiny_conv.cc\n",
        "# Print the source file\n",
        "!cat /content/tiny_conv.cc"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}